import sqlite3

# This module contains three functions:
#   1. save_data - inserts data into the database, and automatically creates tables and
#       organizes the data as necessary.
#   2. retrieve_data - generates SQLite queries to retrieve the requested data from the
#       database and returns the data.
#   3. get_info - reads the database and returns a dictionary that acts as an index to
#       the data in the database. This dictionary is used by the menu function to build
#       out the menu.
# The functions in this module take the following arguments:
#   save_data and retrieve_data take the following arguments:
#   database = database file name in the format "databasename.db". This will be the same 
#       throughout the session.
#   module = the name of the top-level module you're in. Can't have spaces or special
#       characters. Should be EXACTLY matching one of the following:
#       "ReconANDOSINT", "ScanANDEnum", "VulnAnalysis", "Exploitation", or "AuxModules"
#   lvl1, lvl2, lvl3 = the names of the next levels of modules, on down to the one you're 
#       in. If you don't go all the way down to lvl3, input Null. Can have spaces. Should
#       exactly match the name of the module as written in the menu.
#   data = should be the data generated by the module. It will be unchanged when it is
#       passed in and out of the database.

def save_data(database, module, lvl1, lvl2, lvl3, host, data):
    connection = sqlite3.connect(database)
    cursor = connection.cursor()
    check_table = 'SELECT name FROM sqlite_master WHERE type=\'table\' AND name=\'{}\''.format(module)
    cursor.execute(check_table)
    checker = cursor.fetchone()
    if checker == None:
        create_table = 'CREATE TABLE {} (\
            id INTEGER PRIMARY KEY, \
            lvl1 text not null, \
            lvl2 text, \
            lvl3 text, \
            number int, \
            data blob);'.format(module)
        cursor.execute(create_table)
        connection.commit()
    check_entries = 'SELECT number FROM {} WHERE lvl1="{}" AND lvl2="{}" AND lvl3="{}"'.format(module, lvl1, lvl2, lvl3)
    cursor.execute(check_entries)
    printable = cursor.fetchall()
    ctr = 0
    high = []
    while ctr < len(printable):
        high.append(printable[ctr][0])
        ctr += 1
    if high == []:
        highest = 1
    else:
        highest = max(high) + 1
    insert_into_table = 'INSERT INTO {} (lvl1, lvl2, lvl3, number, data) VALUES (?, ?, ?, ?, ?);'.format(module)
    table_entry_parameters = (lvl1, lvl2, lvl3, highest, data)
    cursor.execute(insert_into_table, table_entry_parameters)
    connection.commit()
    connection.close()

def retrieve_data(database, module, lvl1, lvl2, lvl3, num):
    connection = sqlite3.connect(database)
    cursor = connection.cursor()
    if lvl3 == "":
        if lvl2 == "":
            select_from_table = 'SELECT data FROM {} WHERE lvl1="{}" AND lvl2="" AND lvl3="" AND number={}'.format(module, lvl1, num)
        else:
            select_from_table = 'SELECT data FROM {} WHERE lvl1="{}" AND lvl2="{}" AND lvl3="" AND number={}'.format(module, lvl1, lvl2, num)
    else:
        select_from_table = 'SELECT data FROM {} WHERE lvl1="{}" AND lvl2="{}" AND lvl3="{}" AND number={}'.format(module, lvl1, lvl2, lvl3, num)
    cursor.execute(select_from_table)
    returned_data = cursor.fetchall()
    connection.close()
    return returned_data[0][0]

# This function receives only the file name of the database. It reads what is in the 
# database and returns a dictionary that acts as an index to the data saved in the
# database. This dictionary is read by the menu function to build out the menu.
def get_info(database):
    modules = ['ReconANDOSINT', 'ScanANDEnum', 'VulnAnalysis', 'Exploitation', 'AuxModules']
    saved_modules = []
    data_saved = {}
    connection = sqlite3.connect(database)
    cursor = connection.cursor()
    for module in modules:
        check_table = 'SELECT name FROM sqlite_master WHERE type=\'table\' AND name=\'{}\''.format(module)
        cursor.execute(check_table)
        check = cursor.fetchone()
        if check != None:
            saved_modules.append(module)
    for module in saved_modules:
        if module not in data_saved.keys():
            data_saved[module] = {}
        ctr = 1
        while True:
            cmd = 'SELECT lvl1, lvl2, lvl3, number FROM {} WHERE id={}'.format(module, ctr)
            cursor.execute(cmd)
            grab_data = cursor.fetchone()
            if grab_data == None:
                break
            else:
                if grab_data[0] not in data_saved[module].keys():
                    data_saved[module][grab_data[0]] = {}
                if grab_data[1] != '':
                    if grab_data[1] not in data_saved[module][grab_data[0]].keys():
                        data_saved[module][grab_data[0]][grab_data[1]] = {}
                if grab_data[2] == '':
                    if grab_data[1] == '':
                        data_saved[module][grab_data[0]].update( {grab_data[3] : "Data"} )
                    else:
                        data_saved[module][grab_data[0]][grab_data[1]].update( {grab_data[3] : "Data"} )
                else:
                    data_saved[module][grab_data[0]][grab_data[1]][grab_data[2]].update( {grab_data[3] : "Data"} )
            ctr += 1
    connection.close()
    return data_saved